// @Date         2021-12-29
// @Author       Aaron Zhang
// @Description   
// @FileName      tracer.go   

package utils

import (
	"context"
	"github.com/gin-gonic/gin"
	"github.com/rs/xid"
	"net/http"
)

type Tracer struct {
	TraceId string
}

func NewTracer() *Tracer {
	return &Tracer{}
}

func TraceHttpReq(req *http.Request) *Tracer {
	tc := NewTracer()

	if req.Header.Get("tracerId") == "" {
		traceId := xid.New().String()
		tc.TraceId = traceId
	}

	return tc
}

func (t *Tracer) InjectToHttp(r *http.Request) *http.Request {
	r.Header.Set("tracerId", t.TraceId)

	return r
}

func (t *Tracer) InjectToContext(trace *Tracer, ctx context.Context) context.Context {
	if ginCtx, ok := ctx.(*gin.Context); ok {
		ginCtx.Set("trace", t)
		return ginCtx
	} else {
		ctx := context.WithValue(ctx, "trace", t)
		return ctx
	}
}
